import json
import ast

from fit_py_ask_number_demo.utils import convert_str_to_json
from fit_py_ask_number_demo.configs.prompts import dsl_generate_prompt
from fit_py_ask_number_demo.client import llm_client
from fit_py_ask_number_demo.data_managers.query_model_manager import QueryModelManager


class DslGenerate:
    @staticmethod
    def generate_complete_query(query, key_info, field_info):
        time_info = f"""问题的年份为{key_info.get("year")}，月份为{key_info.get("month")}"""
        complete_query = f"用户问题：“{query}”，{time_info}，问题中关键词与数据库字段的映射关系如下："
        for key, value in field_info.items():
            complete_query += f"“{key}”表示{value};"
        complete_query = complete_query[:-1] + "。"
        return complete_query

    @staticmethod
    def get_dsl_example(query_type):
        dsl_example_str = "<examples>"
        dsl_examples = QueryModelManager().get_dsl_example_by_query_type(query_type)
        dsl_examples = ast.literal_eval(dsl_examples)
        for key, value in dsl_examples.items():
            dsl_example_str += f"\ninput: {key} \noutput: {value}\n"
        return dsl_example_str


def get_complete_query_impl(query, key_info, field_info):
    key_info = convert_str_to_json(key_info)
    field_info = convert_str_to_json(field_info)
    complete_query = DslGenerate().generate_complete_query(query, key_info, field_info)
    return complete_query


def get_dsl_example_impl(query_type):
    dsl_example = DslGenerate().get_dsl_example(query_type)
    return dsl_example


def dsl_generate_impl(complete_query, query_type):
    dsl_example = DslGenerate().get_dsl_example(query_type)
    system_content = dsl_generate_prompt + f"<Examples>\n{dsl_example}"
    user_content = f"<Question>\n{complete_query}"
    dsl = llm_client.deepseek_client(system_content, user_content)
    return dsl


if __name__ == "__main__":
    query = "24年战略地区东进计划产互KPI收入"
    query_type = "simple"
    key_info = "{\"year\": [2024], \"month\": None, \"rank\": False, \"comparison\": False}"
    field_info = """{"产互KPI": {"product_lv1": {"in": ["产互KPI"]}}, "北京": {"province": {"in": ['北京']}}}"""
    comp = get_complete_query_impl(query, key_info, field_info)
    dsl = dsl_generate_impl(comp, query_type)
    print(comp)
    print(dsl)
